import * as md from '@babel/traverse';
import _ from 'lodash';

const removeStar = (str) => {
  return str
  .replace(/^\*\x20*/, '') // 开头的*号
  .replace(/\r?\n\x20*\*\x20+/g, '\n') // 每行开头的*号
};

const traverse = md.default.default || md.default;

/**
 * 从AST中提取React箭头函数组件及其leading comments
 * @param {Object} ast - 输入的AST
 * @returns {Map<string, string>} - 返回组件名称到注释的映射
 */
export default function extractReactComponents(ast) {
  const componentMap = new Map();
  
  traverse(ast, {
    FunctionDeclaration(path) {
      const { node } = path;
      
      // 检查是否是React组件（返回JSX）
      if (_.get(node, 'body.body[0].type') === 'ReturnStatement' && 
          _.get(node, 'body.body[0].argument.type') === 'JSXElement') {
        // 获取组件名称和注释
        const componentName = _.get(node, 'id.name');
        const comments = _.chain(node.leadingComments)
          .filter({ type: 'CommentBlock' })
          .map(c => removeStar( _.trim(c.value)))
          .join('\n')
          .value();
        
        if (componentName) {
          componentMap.set(componentName, comments);
        }
      }
    },
    ClassDeclaration(path) {
      const { node } = path;
      
      // 检查是否是React类组件（有render方法返回JSX）
      const renderMethod = node.body.body.find(
        method => method.type === 'ClassMethod' && method.key.name === 'render'
      );
      
      if (renderMethod && 
          _.get(renderMethod, 'body.body[0].type') === 'ReturnStatement' &&
          _.get(renderMethod, 'body.body[0].argument.type') === 'JSXElement') {
        // 获取组件名称和注释
        const componentName = _.get(node, 'id.name');
        const comments = _.chain(node.leadingComments)
          .filter({ type: 'CommentBlock' })
          .map(c => removeStar( _.trim(c.value)))
          .join('\n')
          .value();
        
        if (componentName) {
          componentMap.set(componentName, comments);
        }
      }
    },
    VariableDeclaration(path) {
      const { node } = path;
      
      // 检查变量声明是否包含箭头函数或函数表达式
      node.declarations.forEach(declarator => {
        if (declarator.init && 
            (declarator.init.type === 'ArrowFunctionExpression' || 
             declarator.init.type === 'FunctionExpression')) {
          const initNode = declarator.init;
          
          // 检查是否是React组件（返回JSX）
          if (_.get(initNode, 'body.type') === 'JSXElement' || 
              (_.get(initNode, 'body.type') === 'BlockStatement' && 
               _.get(initNode, 'body.body[0].type') === 'ReturnStatement' &&
               _.get(initNode, 'body.body[0].argument.type') === 'JSXElement')) {
            
            // 获取组件名称
            const componentName = _.get(declarator, 'id.name');
            
            // 获取注释
            const comments = _.chain(node.leadingComments)
              .filter({ type: 'CommentBlock' })
              .map(c => removeStar( _.trim(c.value)))
              .join('\n')
              .value();
            
            if (componentName) {
              componentMap.set(componentName, comments);
            }
          }
        }
      });
    }
  });
  
  return componentMap;
}